#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
修复PostgreSQL数据库表结构的脚本
确保telegrams表结构与pg_models.py中定义的一致
"""

import psycopg2
from config import PG_HOST, PG_PORT, PG_DATABASE, PG_USER, PG_PASSWORD

def fix_postgres_schema():
    """修复PostgreSQL数据库表结构"""
    try:
        # 连接到PostgreSQL数据库
        print("尝试连接到PostgreSQL数据库...")
        conn = psycopg2.connect(
            host=PG_HOST,
            port=PG_PORT,
            database=PG_DATABASE,
            user=PG_USER,
            password=PG_PASSWORD
        )
        conn.autocommit = True  # 设置自动提交，确保修改生效
        cursor = conn.cursor()
        print("成功连接到PostgreSQL数据库")
        
        # 检查并修复表结构
        print("开始检查并修复表结构...")
        
        # 检查telegram_id字段是否存在，如果不存在则添加
        cursor.execute("""
        SELECT EXISTS (
            SELECT 1 
            FROM information_schema.columns 
            WHERE table_name='telegrams' AND column_name='telegram_id'
        )
        """)
        has_telegram_id = cursor.fetchone()[0]
        
        if not has_telegram_id:
            print("添加telegram_id字段...")
            cursor.execute("ALTER TABLE telegrams ADD COLUMN telegram_id INTEGER")
            print("telegram_id字段添加成功")
            
            # 将id_original的值复制到telegram_id字段
            print("将id_original的值复制到telegram_id字段...")
            cursor.execute("UPDATE telegrams SET telegram_id = id_original")
            print(f"已更新 {cursor.rowcount} 条记录的telegram_id字段")
        else:
            print("telegram_id字段已存在")
        
        # 检查mobile_url字段是否存在，如果不存在则添加
        cursor.execute("""
        SELECT EXISTS (
            SELECT 1 
            FROM information_schema.columns 
            WHERE table_name='telegrams' AND column_name='mobile_url'
        )
        """)
        has_mobile_url = cursor.fetchone()[0]
        
        if not has_mobile_url:
            print("添加mobile_url字段...")
            cursor.execute("ALTER TABLE telegrams ADD COLUMN mobile_url TEXT")
            print("mobile_url字段添加成功")
        else:
            print("mobile_url字段已存在")
        
        # 检查audio_url字段是否存在，如果不存在则添加
        cursor.execute("""
        SELECT EXISTS (
            SELECT 1 
            FROM information_schema.columns 
            WHERE table_name='telegrams' AND column_name='audio_url'
        )
        """)
        has_audio_url = cursor.fetchone()[0]
        
        if not has_audio_url:
            print("添加audio_url字段...")
            cursor.execute("ALTER TABLE telegrams ADD COLUMN audio_url TEXT")
            print("audio_url字段添加成功")
        else:
            print("audio_url字段已存在")
        
        # 检查created_at字段是否存在，如果不存在则添加
        cursor.execute("""
        SELECT EXISTS (
            SELECT 1 
            FROM information_schema.columns 
            WHERE table_name='telegrams' AND column_name='created_at'
        )
        """)
        has_created_at = cursor.fetchone()[0]
        
        if not has_created_at:
            print("添加created_at字段...")
            cursor.execute("ALTER TABLE telegrams ADD COLUMN created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP")
            print("created_at字段添加成功")
        else:
            print("created_at字段已存在")
        
        # 添加缺失的索引
        print("检查并创建缺失的索引...")
        
        # 检查idx_telegram_id索引是否存在
        cursor.execute("""
        SELECT EXISTS (
            SELECT 1 
            FROM pg_indexes 
            WHERE tablename='telegrams' AND indexname='idx_telegram_id'
        )
        """)
        has_idx_telegram_id = cursor.fetchone()[0]
        
        if not has_idx_telegram_id:
            print("创建idx_telegram_id索引...")
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_telegram_id ON telegrams(telegram_id)")
            print("idx_telegram_id索引创建成功")
        else:
            print("idx_telegram_id索引已存在")
        
        # 检查idx_time索引是否存在
        cursor.execute("""
        SELECT EXISTS (
            SELECT 1 
            FROM pg_indexes 
            WHERE tablename='telegrams' AND indexname='idx_time'
        )
        """)
        has_idx_time = cursor.fetchone()[0]
        
        if not has_idx_time:
            print("创建idx_time索引...")
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_time ON telegrams(time)")
            print("idx_time索引创建成功")
        else:
            print("idx_time索引已存在")
        
        print("表结构修复完成！")
        
        # 关闭连接
        cursor.close()
        conn.close()
        
        return True
    except Exception as e:
        print(f"修复表结构时出错: {e}")
        return False

if __name__ == "__main__":
    print("=== 开始修复PostgreSQL数据库表结构 ===")
    fix_postgres_schema()